# ''' Python 代码开始

import cv2
import numpy as np
import glob
import yaml
import os
import argparse
import abc

# ===================================================================
# 辅助函数: 用于加载和转换数据
# ===================================================================

def load_transformation_matrices(path_pattern):
    """
    从匹配路径模式的文件中加载一系列4x4变换矩阵.
    每个文件应包含一个用空格分隔的4x4矩阵.
    """
    matrices = []
    file_list = sorted(glob.glob(path_pattern))
    if not file_list:
        raise FileNotFoundError(f"错误: 在路径 '{path_pattern}' 未找到任何位姿文件.")
        
    for fname in file_list:
        try:
            matrix = np.loadtxt(fname)
            if matrix.shape != (4, 4):
                raise ValueError(f"文件 '{fname}' 中的矩阵形状不是4x4.")
            matrices.append(matrix)
        except Exception as e:
            print(f"WARNING: 无法加载或解析文件 {fname}: {e}")
            continue
    return matrices

def separate_R_t(matrices):
    """将4x4变换矩阵列表分离为旋转矩阵和位移向量列表."""
    Rs = [m[:3, :3] for m in matrices]
    ts = [m[:3, 3] for m in matrices]
    return Rs, ts

# ===================================================================
# 策略模式: 定义手眼标定求解器的基类和具体实现
# ===================================================================

class HandEyeSolver(abc.ABC):
    """
    手眼标定求解器的抽象基类 (接口).
    """
    def __init__(self, method_name):
        self.method_name = method_name

    @abc.abstractmethod
    def solve(self, R_gripper2base, t_gripper2base, R_target2cam, t_target2cam):
        """
        执行手眼标定求解过程.

        参数:
        - R_gripper2base: 从机器人基座到末端执行器的旋转矩阵列表.
        - t_gripper2base: 从机器人基座到末端执行器的平移向量列表.
        - R_target2cam: 从相机到标定板的旋转矩阵列表.
        - t_target2cam: 从相机到标定板的平移向量列表.

        返回:
        - R_cam_to_end: 计算出的旋转矩阵.
        - t_cam_to_end: 计算出的平移向量.
        """
        pass

class OpenCVHandEyeSolver(HandEyeSolver):
    """使用OpenCV内建函数的手眼标定求解器."""
    
    METHOD_FLAGS = {
        'TSAI': cv2.CALIB_HAND_EYE_TSAI,
        'PARK': cv2.CALIB_HAND_EYE_PARK,
        'HORAUD': cv2.CALIB_HAND_EYE_HORAUD,
        'ANDREFF': cv2.CALIB_HAND_EYE_ANDREFF,
        'DANIILIDIS': cv2.CALIB_HAND_EYE_DANIILIDIS
    }

    def __init__(self, method_name):
        super().__init__(method_name)
        self.method_flag = self.METHOD_FLAGS.get(method_name.upper())
        if self.method_flag is None:
            raise ValueError(f"错误: 不支持的OpenCV标定方法 '{method_name}'. "
                             f"支持的方法有: {list(self.METHOD_FLAGS.keys())}")

    def solve(self, R_gripper2base, t_gripper2base, R_target2cam, t_target2cam):
        """使用cv2.calibrateHandEye执行求解."""
        print(f"INFO: 使用OpenCV的 '{self.method_name}' 方法进行求解...")
        R_cam_to_end, t_cam_to_end = cv2.calibrateHandEye(
            R_gripper2base, t_gripper2base,
            R_target2cam, t_target2cam,
            method=self.method_flag
        )
        return R_cam_to_end, t_cam_to_end

# ===================================================================
# 主标定类: HandEyeCalibrator
# ===================================================================

class HandEyeCalibrator:
    """
    一个功能全面的手眼标定类 (AX=XB).

    核心功能:
    - 通过YAML配置文件驱动，易于配置和使用。
    - 采用策略模式，轻松支持多种标定算法 (Tsai-Lenz, Park-Martin等)。
    - 支持 "Eye-in-Hand" 和 "Eye-to-Hand" 两种配置。
    - 从文件中加载机器人位姿和相机位姿数据。
    - 执行手眼标定，计算相机和机器人末端（或基座）之间的变换。
    - 将最终的变换矩阵保存为npz或yaml文件。
    """
    
    SOLVER_HANDLERS = {
        'TSAI': OpenCVHandEyeSolver,
        'PARK': OpenCVHandEyeSolver,
        'HORAUD': OpenCVHandEyeSolver,
        'ANDREFF': OpenCVHandEyeSolver,
        'DANIILIDIS': OpenCVHandEyeSolver
    }

    def __init__(self, config_path):
        """
        初始化, 加载配置并创建对应的求解器.
        :param config_path: 配置文件 (hand_eye_config.yaml) 的路径.
        """
        print("INFO: 初始化手眼标定器...")
        with open(config_path, 'r') as f:
            self.config = yaml.safe_load(f)

        self.robot_poses_path = self.config['robot_poses_path']
        self.camera_poses_path = self.config['camera_poses_path']
        self.output_path = self.config['output_path']
        self.save_format = self.config['save_format']
        
        self.method = self.config['hand_eye_method']
        self.setup = self.config['calibration_setup'] # 'EYE_IN_HAND' or 'EYE_TO_HAND'

        handler_class = self.SOLVER_HANDLERS.get(self.method.upper())
        if not handler_class:
            raise ValueError(f"错误: 不支持的标定方法 '{self.method}'. "
                             f"支持的类型有: {list(self.SOLVER_HANDLERS.keys())}")
        self.solver = handler_class(self.method)
        print(f"INFO: 已选择标定方法: '{self.method}'")
        print(f"INFO: 标定设置: '{self.setup}'")

        if not os.path.exists(self.output_path):
            os.makedirs(self.output_path)
            print(f"INFO: 已创建输出目录: {self.output_path}")

    def calibrate(self):
        """
        执行手眼标定主流程.
        """
        print("\nINFO: 开始标定流程...")
        
        # 1. 加载所有位姿数据
        print("INFO: 正在加载机器人和相机的位姿数据...")
        robot_matrices = load_transformation_matrices(self.robot_poses_path)
        camera_matrices = load_transformation_matrices(self.camera_poses_path)

        if len(robot_matrices) != len(camera_matrices):
            raise ValueError(f"错误: 机器人位姿数量 ({len(robot_matrices)}) 和 "
                             f"相机位姿数量 ({len(camera_matrices)}) 不匹配.")
        
        if len(robot_matrices) < 3:
            print(f"WARNING: 位姿数据对少于3个 ({len(robot_matrices)}个), 标定结果可能不可靠.")

        print(f"INFO: 成功加载 {len(robot_matrices)} 对位姿数据.")

        # 2. 根据标定设置准备数据 (AX=XB)
        # cv2.calibrateHandEye 解决的是 AX=XB 问题, 其中:
        # X: gripper -> camera (眼在手上) 或 base -> camera (眼在手外)
        # A: gripper_i -> gripper_j
        # B: camera_i -> camera_j
        # 函数输入需要的是 R_gripper2base 和 R_target2cam
        
        if self.setup == 'EYE_IN_HAND':
            # A: base -> gripper, B: target -> cam, X: gripper -> cam
            # 输入: R_gripper2base, t_gripper2base, R_target2cam, t_target2cam
            R_gripper2base, t_gripper2base = separate_R_t(robot_matrices)
            R_target2cam, t_target2cam = separate_R_t(camera_matrices)
            
        elif self.setup == 'EYE_TO_HAND':
            # A: base -> gripper, B: target -> cam, Y: base -> cam
            # 转换为 AX=XB 形式: inv(A)Y = YB', 其中 A是gripper->base, Y是base->cam, B'是target->cam
            # OpenCV求解器需要: inv(R_base2gripper) 和 R_target2cam.
            # inv(base->gripper) == gripper->base.
            print("INFO: 'EYE_TO_HAND' 设置, 对机器人位姿进行逆变换以适配求解器.")
            robot_matrices_inv = [np.linalg.inv(m) for m in robot_matrices]
            R_gripper2base, t_gripper2base = separate_R_t(robot_matrices_inv) # 这里实际上是 gripper -> base
            R_target2cam, t_target2cam = separate_R_t(camera_matrices)
            
        else:
            raise ValueError(f"错误: 无效的 'calibration_setup': {self.setup}. "
                             "请选择 'EYE_IN_HAND' 或 'EYE_TO_HAND'.")

        # 3. 调用策略进行求解
        R, t = self.solver.solve(R_gripper2base, t_gripper2base, R_target2cam, t_target2cam)

        # 4. 整合并保存结果
        # 创建4x4齐次变换矩阵
        H = np.identity(4)
        H[:3, :3] = R
        H[:3, 3] = t.flatten()
        
        self._save_results(H)
        
    def _save_results(self, H_matrix):
        """保存最终的变换矩阵."""
        print("\n----------- 手眼标定结果 -----------")
        if self.setup == 'EYE_IN_HAND':
            print("相机相对于末端执行器的变换矩阵 (cam_H_gripper):")
        else: # EYE_TO_HAND
            print("相机相对于机器人基座的变换矩阵 (cam_H_base):")
        
        np.set_printoptions(precision=4, suppress=True)
        print(H_matrix)
        print("------------------------------------\n")

        if self.save_format == 'npz':
            save_path = os.path.join(self.output_path, 'hand_eye_result.npz')
            np.savez(save_path, H=H_matrix)
        elif self.save_format == 'yaml':
            save_path = os.path.join(self.output_path, 'hand_eye_result.yaml')
            data = {'hand_eye_transformation': H_matrix.tolist()}
            with open(save_path, 'w') as f:
                yaml.dump(data, f, default_flow_style=False, sort_keys=False)
        else:
            print(f"WARNING: 不支持的保存格式 '{self.save_format}'. 结果未保存.")
            return

        print(f"INFO: 标定结果已保存至: {save_path}")
        print("\n手眼标定流程完成！")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="手眼标定脚本 (支持 'Eye-in-Hand' 和 'Eye-to-Hand')")
    parser.add_argument('--config', type=str, default='hand_eye_config.yaml', 
                        help='配置文件的路径 (默认为: hand_eye_config.yaml)')
    args = parser.parse_args()
    
    if not os.path.exists(args.config):
        print(f"ERROR: 配置文件 '{args.config}' 不存在. 请确保该文件存在或通过 --config 参数指定正确路径。")
    else:
        try:
            calibrator = HandEyeCalibrator(args.config)
            calibrator.calibrate()
        except (ValueError, FileNotFoundError, KeyError) as e:
            print(f"\nERROR: 处理失败. 可能是配置文件内容有误或文件路径问题。")
            print(f"  > 错误详情: {e}")

# ''' Python代码结束